import os

import pandas as pd

root_path = r'H:\notebook\LeNet\note1_extra_selfdata\hotdog'
train_f = r'train'
test_f = r'test'
hotdog_f = r'hotdog'
nhotdog_f = r'not-hotdog'


def main():
    train_arr = []

    tmp_path = os.path.join(root_path, train_f, hotdog_f)
    tmp = os.listdir(tmp_path)
    for fn in tmp:
        train_arr.append([hotdog_f + "\\" + fn, 1])

    tmp_path = os.path.join(root_path, train_f, nhotdog_f)
    tmp = os.listdir(tmp_path)
    for fn in tmp:
        train_arr.append([nhotdog_f + "\\" + fn, 0])

    test_arr = []
    tmp_path = os.path.join(root_path, test_f, hotdog_f)
    tmp = os.listdir(tmp_path)
    for fn in tmp:
        test_arr.append([hotdog_f + "\\" + fn, 1])

    tmp_path = os.path.join(root_path, test_f, nhotdog_f)
    tmp = os.listdir(tmp_path)
    for fn in tmp:
        test_arr.append([nhotdog_f + "\\" + fn, 0])

    train_df = pd.DataFrame(train_arr, columns=['path', 'label'])
    test_df = pd.DataFrame(test_arr, columns=['path', 'label'])
    train_df.to_csv(os.path.join(root_path, "train.csv"))
    test_df.to_csv(os.path.join(root_path, "test.csv"))


if __name__ == "__main__":
    main()
